import os
import torch
import torch.nn as nn

from bert_chinese_encode import get_bert_encode_for_single
from rnn_model import RNN


# 预加载的模型参数路径
MODEL_PATH = './BERT_RNN_22000.pth'

# 隐层节点数，输入层尺寸，类别数都和训练时相同即可
n_hidden = 128
input_size = 768
n_categories = 2

# 实例化RNN模型，并加载保存模型参数
rnn = RNN(input_size, n_hidden, n_categories)
rnn.load_state_dict(torch.load(MODEL_PATH))

def _test(text_tensor):
    """
    计算预测结果函数
    :param text_tensor:
    :return:
    """
    hidden = rnn.init_hidden()
    with torch.no_grad():
        for i in range(text_tensor.size()[1]):
            output, hidden = rnn(text_tensor[0][i].unsqueeze(0), hidden)

    return output

def predict(input_text):
    """
    把输入的症状的文本进行计算得到预测标签
    :param input_text: 症状的文本
    :return: 标签
    """
    with torch.no_grad():
        output = _test(get_bert_encode_for_single(input_text))

        _, topi = output.topk(1, 1)
        return topi.item()

def batch_predict(input_path, output_path):
    """
    批量审核每个疾病的症状
    :param input_path: 待审核疾病文件所在的目录
    :param output_path: 审核完成后疾病文件所在的目录
    :return:
    """

    csv_list = os.listdir(input_path)
    for disease in csv_list:
        print(disease)
        with open(os.path.join(input_path, disease), 'r') as fr:
            with open(os.path.join(output_path, disease), 'w') as fw:
                lines = fr.readlines()
                for line in lines:
                    result = predict(line)
                    if result:
                        fw.write(line)
                    else:
                        pass


if __name__ == '__main__':
    # print(predict("点瘀样尖针性发多"))
    input_path = "/root/ai17bj/doctor_offline/structured/noreview/"
    output_path = "/root/ai17bj/doctor_offline/structured/reviewed/"
    batch_predict(input_path, output_path)
